#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File        : mngs_parse_jsonl.py
@Author      : Bing Liang
@Email       : believer19940901@gmail.com
@Date        : 2025/11/03 16:00
@Description : 解析 NCBI datasets 下载的基因组元文件，拼接静态 FTP 下载地址
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
import json


def acc_to_url(acc: str, asm_name: str) -> str:
    """将 GCF/GCA 编号和组装名拼接为 NCBI FTP 静态路径"""
    prefix = acc.split('_')[0]
    num = acc.split('_')[1].split('.')[0]
    parts = [num[i:i + 3] for i in range(0, len(num), 3)]
    return f"https://ftp.ncbi.nlm.nih.gov/genomes/all/{prefix}/{'/'.join(parts)}/{acc}_{asm_name}"


def mngs_parse_jsonl(args: Namespace) -> None:
    """解析 JSONL 文件并生成对应的 genome 下载链接"""
    in_file = args.in_file.resolve()
    out_file = args.out_file.resolve()
    out_file.parent.mkdir(parents=True, exist_ok=True)

    with in_file.open("r", encoding="utf-8") as fr, out_file.open("w", encoding="utf-8") as fw:
        count = 0
        for line in fr:
            line = line.strip()
            if not line:
                continue

            try:
                dic = json.loads(line)
                acc = dic.get("currentAccession")
                asm = dic.get("assemblyInfo", {}).get("assemblyName")

                if not acc or not asm:
                    continue

                base_url = acc_to_url(acc, asm)
                fna_url = f"{base_url}/{acc}_{asm}_genomic.fna.gz"
                fw.write(f"{fna_url}\n")
                count += 1

            except json.JSONDecodeError:
                continue

    print(f"已生成 {count} 条下载链接，输出文件：{out_file}")


def main():
    parser = ArgumentParser(description="解析 dataset 基因组元文件，拼接 NCBI 静态下载链接")
    parser.add_argument("--in_file", type=Path, required=True, help="输入 JSONL 文件路径")
    parser.add_argument("--out_file", type=Path, required=True, help="输出下载链接文件路径")
    args = parser.parse_args()
    mngs_parse_jsonl(args)


if __name__ == "__main__":
    main()
